#!/usr/bin/env python2.7
#-*- coding:utf-8 -*-
#Author='He Rensheng'
#Email='hrs323@126.com'

# Some standard imports
import io, os
import numpy as np
import torch.onnx
from torch.autograd import Variable
# import some class
from SqueezeNet import squeezenet1_1
from utils import *


# Get trained squeezenet model
trained_model_savedir = 'pre-trained-model'
torch_model = squeezenet1_1(trained_model_savedir, True)

# set the save path of the exported onnx model
if not os.path.exists('output/onnx'):
    os.makedirs('output/onnx')
onnx_path = 'output/onnx/squeezenet.onnx'

# set the save path of the exported caffe2 model
if not os.path.exists('output/caffe2'):
    os.makedirs('output/caffe2')
caffe2_init_path = 'output/caffe2/squeeze_init_net.pb'
caffe2_predict_path = 'output/caffe2/squeeze_predict_net.pb'

# export the pytorch model as onnx model
torch_out, x = pytorch2onnx(torch_model, onnx_path)

# export onnx model to caffe2 model and verify
onnx2caffe2(onnx_path, torch_out, x, caffe2_init_path, caffe2_predict_path)
